package main

import (
	"fmt"
	"net/http"
	"time"
)

type PressureGauge struct {
	ch chan struct{}
}

func New(limit int) *PressureGauge {
	ch := make(chan struct{}, limit)
	for i := 0; i < limit; i++ {
		ch <- struct{}{}
	}

	return &PressureGauge{ch: ch}
}

func (g *PressureGauge) process(f func()) error {
	select {
	case <-g.ch:
		f()
		g.ch <- struct{}{}
		return nil
	default:
		return fmt.Errorf("处理数据过多，触发限流")
	}
}

func doSomeThingShouldBeLimited() string {
	time.Sleep(5 * time.Second)
	return "done"
}

func main() {
	pg := New(1)
	http.HandleFunc("/request", func(writer http.ResponseWriter, request *http.Request) {
		err := pg.process(func() {
			writer.Write([]byte(doSomeThingShouldBeLimited()))
		})

		if err != nil {
			writer.WriteHeader(http.StatusTooManyRequests)
			writer.Write([]byte("Too Many Request"))
		}
	})
	err := http.ListenAndServe(":8080", nil)
	if err != nil {
		fmt.Println(err)
	}
}
